package com.mttk.api.addon.ai.action;

import java.util.ArrayList;
import java.util.List;

import org.apache.http.HttpEntity;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import org.bson.Document;

import com.mttk.api.addon.Action;
import com.mttk.api.addon.AdapterConfig;
import com.mttk.api.addon.Pipeline;
import com.mttk.api.addon.RuntimeContext;
import com.mttk.api.addon.UnretryException;
import com.mttk.api.addon.ai.pool.AiPoolService;
import com.mttk.api.addon.ai.pool.model.AiModel;
import com.mttk.api.addon.annotation.ActionFlag;
import com.mttk.api.addon.annotation.ui.Control;
import com.mttk.api.util.StringUtil;

@ActionFlag(key = "aiChat", name = "AI聊天", folder = "aiFolder", i18n = "/com/mttk/api/addon/ai/i18n")
@Control(key = "aiModel", label = "AI模型", mode = "chooser", mandatory = true, props = { "type:pool.ai_model_pool" })
@Control(key = "url", label = "URL", defaultVal = "/chat/completions")
@Control(key = "responseMode", label = "结果DUMP方式", description = "对话结果记录到管道 ai_response_content,完整信息记录到ai_response_full", mandatory = true, mode = "select", defaultVal = "content", props = "options:none:不,content:记录对话结果,both:记录对话结果和完整信息")
@Control(key = "chat_session", label = "聊天会话", description = "相同会话值的聊天会自动获取聊天历史")


@Control(key = "prompt_system", size = 1, label = "系统提示词", props = { "type:multiple", "rows:5" })
@Control(key = "prompt_user", size = 1, label = "用户提示词", mandatory = true, props = { "type:multiple", "rows:5" })
public class AiChatAction implements Action {

	@Override
	public Object invoke(RuntimeContext c, Pipeline p, AdapterConfig a) throws UnretryException, Exception {
		// AI model
		AiModel aiModel = c.getServer().getService(AiPoolService.class).getPooledObject(a.getStringMandatory("aiModel"))
				.get();
		//

		try (CloseableHttpClient httpClient = buildHttpClient()) {
			String url = aiModel.getUrl() + a.getString("url", "");
			HttpPost httpPost = new HttpPost(url);
			// Request
			Document requestBody = buildRequestBody(c, p, a, aiModel);
			String bodyStr = requestBody.toJson();

			c.getLogger().debug("Request body:\n" + bodyStr);
			//
			HttpEntity entity = new StringEntity(bodyStr, ContentType.APPLICATION_JSON);
			httpPost.setEntity(entity);
			//
			//
			httpPost.addHeader("Content-Type", "application/json");
			httpPost.addHeader("Accept", "application/json");
			httpPost.addHeader("Authorization", "Bearer " + aiModel.getConfig().getStringMandatory("appKey"));
			//
			try (CloseableHttpResponse resp = httpClient.execute(httpPost)) {
			
				//
				HttpEntity respBody = resp.getEntity();
				String respStr = EntityUtils.toString(respBody);
				//
				c.getLogger().debug("Response:\n" + resp);
				c.getLogger().debug("Response body:\n" + respStr);

				//
				return handleResponse(c, p, a, aiModel, requestBody, resp, respStr);
			}
		}
	}

	private CloseableHttpClient buildHttpClient() {
		HttpClientBuilder httpClientBuilder = HttpClients.custom();
		// 超时设置
		RequestConfig.Builder requestConfigBuilder = RequestConfig.custom();
		requestConfigBuilder.setConnectTimeout(60 * 1000);
		requestConfigBuilder.setSocketTimeout(180 * 1000);

		httpClientBuilder.setDefaultRequestConfig(requestConfigBuilder.build());
		//
		return  httpClientBuilder.build();
	}

	 private Document buildRequestBody(RuntimeContext c, Pipeline p, AdapterConfig a, AiModel aiModel)
			throws UnretryException, Exception {
		Document body = new Document();
		// model
		body.append("model", aiModel.getModel());
		// properties
		aiModel.applyProps(body);
		//ext_props
		handleExtProps(body,aiModel.getConfig());
		// messages
		List<Document> messages = new ArrayList<>();
		body.append("messages", messages);
		//
		String chat_session=a.getString("chat_session");
		if (StringUtil.notEmpty(chat_session)) {
			String chatHistory = p.getString("ai_chat_history_"+chat_session);
			if (StringUtil.notEmpty(chatHistory)) {
				Document chatHistoryDoc = Document.parse(chatHistory);
				List<Document> ms = (List<Document>) chatHistoryDoc.get("messages");
				if (ms != null && ms.size() > 0) {
					//
					for (Document m : ms) {
						messages.add(m);
					}
				}
			}
		}
		//
		String prompt_system = a.getString("prompt_system");
		if (StringUtil.notEmpty(prompt_system)) {
			messages.add(new Document().append("role", "system").append("content", prompt_system));
		}
		String prompt_user = a.getStringMandatory("prompt_user");
		if (StringUtil.notEmpty(prompt_user)) {
			messages.add(new Document().append("role", "user").append("content", prompt_user));
		}
		//
		return body;

	}
	 private void handleExtProps(Document body,AdapterConfig a) {
			String extProps=a.getString("ext_props");

			if(StringUtil.isEmpty(extProps)) {
				return;
			}
			//
			Document d=Document.parse(extProps);
			for(String key:d.keySet()) {
				Object value=d.get(key);

				if(StringUtil.isEmpty(value)) {
					//试图删除
					body.remove(key);
				}else {
					body.put(key, value);
				}
			}
	 }
	private Object handleResponse(RuntimeContext c, Pipeline p, AdapterConfig a, AiModel aiModel, Document requestBody,
			CloseableHttpResponse resp, String respStr) throws UnretryException, Exception {
		//
		String responseMode = a.getString("responseMode", "content");
		if ("both".equals(responseMode)) {
			p.put("ai_response_full", respStr);
		}
		//
		if(resp.getStatusLine().getStatusCode()!=200) {
			throw new RuntimeException("AI response error. Status :"+resp.getStatusLine()+",content:"+respStr);
		}
		//
		Document respDoc = null;
		try {
		respDoc=Document.parse(respStr);
		}catch(Exception e) {			
			throw new RuntimeException("AI response error. Status :"+resp.getStatusLine()+",content:"+respStr);
		}
		if(respDoc.containsKey("error")) {
			throw new RuntimeException("AI response error. Status :"+resp.getStatusLine()+",content:"+respStr);
		}
		//检查错误信息
//		System.out.println("respStr="+respStr);
		// 获取返回内容
		Object choicesObj=respDoc.get("choices");
		if(choicesObj==null) {
			return null;
		}
		List<Document> choices = (List<Document>)choicesObj ;
		if (choices.size() > 0) {

			// 获取最后一条choice
			Document choice = choices.get(choices.size() - 1);
			Document message = (Document) choice.get("message");
			Object content = message.get("content");
			String contentStr = null;
			if (content == null) {
				contentStr = "";
			} else if (content instanceof Document) {
				contentStr = ((Document) content).toJson();
			} else {
				contentStr = content.toString();
			}
			//
			if (!"none".equals(responseMode)) {
				p.put("ai_response_content", contentStr);
			}
		}
		//
		String chat_session=a.getString("chat_session");
		if (StringUtil.notEmpty(chat_session)) {
			handleRecordChatHistory(chat_session,p, aiModel, requestBody, respDoc);
		}
		// 返回值
		respDoc.remove("choices");
		return respDoc.toJson();
	}

	private void handleRecordChatHistory(String chat_session,Pipeline p, AiModel aiModel, Document requestBody, Document respDoc) {
		List<Document> list = new ArrayList<>();
		// 请求里的消息
		List<Document> messages = (List<Document>) requestBody.get("messages");
		for (Document message : messages) {
			list.add(message);
		}
		// 返回值里的choices
		List<Document> choices = (List<Document>) respDoc.get("choices");
		for (Document choice : choices) {
			list.add((Document)choice.get("message"));
		}
		//
		Document history = new Document("messages", list);
		//
		p.put("ai_chat_history_"+chat_session, history.toJson());

	}
}
